import torch
from torch import nn
from torch.nn import init
import torch.nn.functional as F
import math
from torch.autograd import Variable
import numpy as np

from .deeplab_resnet import resnet50_locate
from .vgg import vgg16_locate

# 统一配置管理
model_configs = {
    'vgg': {
        'convert': [[128, 256, 512, 512, 512], [64, 128, 256, 512, 512]],
        'deep_pool': [[512, 512, 256, 128],
                      [512, 256, 128, 128],
                      [True, True, True, False],
                      [True, True, True, False]],
        'score': 128
    },
    'resnet': {
        'convert': [[64, 256, 512, 1024, 2048], [128, 256, 256, 512, 512]],
        'deep_pool': [[512, 512, 256, 256, 128],
                      [512, 256, 256, 128, 128],
                      [False, True, True, True, False],
                      [True, True, True, True, False]],
        'score': 128
    }
}


class ConvertLayer(nn.Module):
    """特征转换层，用于统一不同主干网络的通道数"""

    def __init__(self, in_channels_list, out_channels_list):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(in_ch, out_ch, 1, bias=False),
                # 移除 inplace 操作
                nn.ReLU()
            )
            for in_ch, out_ch in zip(in_channels_list, out_channels_list)
        ])

    def forward(self, features):
        return [layer(feat) for layer, feat in zip(self.layers, features)]


class DeepPoolLayer(nn.Module):
    """深度池化层，整合多尺度信息"""

    def __init__(self, in_channels, out_channels, upsample, fuse):
        super().__init__()
        self.pool_convs = nn.ModuleList([
            nn.Sequential(
                nn.AvgPool2d(kernel_size=size, stride=size),
                nn.Conv2d(in_channels, in_channels, 3, padding=1, bias=False)
            )
            for size in [2, 4, 8]
        ])

        self.merge_conv = nn.Sequential(
            # 移除 inplace 操作
            nn.ReLU(),
            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False)
        )

        self.fuse_conv = nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False) if fuse else nn.Identity()
        self.upsample = upsample
        self.fuse = fuse

    def forward(self, x, x2=None, x3=None):
        # 多尺度池化与特征融合
        res = x
        for pool_conv in self.pool_convs:
            pooled = pool_conv(x)
            # 避免就地操作，创建新变量
            upsampled = F.interpolate(pooled, size=x.shape[2:], mode='bilinear', align_corners=True)
            res = res + upsampled  # 使用赋值而非就地加法

        # 特征合并
        res = self.merge_conv(res)

        # 上采样（如果需要）
        if self.upsample and x2 is not None:
            res = F.interpolate(res, size=x2.shape[2:], mode='bilinear', align_corners=True)

        # 特征融合（如果需要）
        if self.fuse and x2 is not None and x3 is not None:
            # 避免就地操作
            fused = res + x2 + x3
            res = self.fuse_conv(fused)

        return res


class ScoreLayer(nn.Module):
    """输出评分层，生成最终的分割图"""

    def __init__(self, in_channels):
        super().__init__()
        self.score = nn.Conv2d(in_channels, 1, kernel_size=1)

    def forward(self, x, target_size=None):
        x = self.score(x)
        if target_size is not None:
            x = F.interpolate(x, size=target_size, mode='bilinear', align_corners=True)
        return x


class PoolNet(nn.Module):
    """PoolNet 主网络结构"""

    def __init__(self, base_model_cfg='vgg'):
        super().__init__()
        self.base_model_cfg = base_model_cfg
        cfg = model_configs[base_model_cfg]

        # 主干网络选择
        if base_model_cfg == 'vgg':
            from .vgg import vgg16_locate
            self.base = vgg16_locate()
        else:  # resnet
            from .deeplab_resnet import resnet50_locate
            self.base = resnet50_locate()

        # 特征转换层（仅resnet需要）
        self.convert = ConvertLayer(*cfg['convert']) if base_model_cfg == 'resnet' else None

        # 深度池化层
        self.deep_pool = nn.ModuleList([
            DeepPoolLayer(in_ch, out_ch, upsample, fuse)
            for in_ch, out_ch, upsample, fuse in zip(*cfg['deep_pool'])
        ])

        # 输出层
        self.score = ScoreLayer(cfg['score'])

        # 初始化权重
        self.apply(self._init_weights)  # 修改这里

    @staticmethod
    def _init_weights(m):
        """Kaiming初始化方案"""
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)

    # 修正后的前向传播
    def forward(self, x):
        input_size = x.shape[2:]
        features, infos = self.base(x)
        if self.convert: features = self.convert(features)
        features = features[::-1]

        current = features[0]
        # 按特征数量决定融合次数
        for i in range(len(features) - 1):  # 关键修正
            current = self.deep_pool[i](current, features[i + 1], infos[i])

        # 处理额外的池化层 (ResNet 需要)
        extra_layers = len(self.deep_pool) - (len(features) - 1)
        for i in range(extra_layers):
            current = self.deep_pool[len(features) - 1 + i](current)

        return self.score(current, input_size)


def build_model(base_model_cfg='vgg'):
    """模型构建函数"""
    return PoolNet(base_model_cfg)


def weights_init(m):
    if isinstance(m, nn.Conv2d):
        m.weight.data.normal_(0, 0.01)
        if m.bias is not None:
            m.bias.data.zero_()